package com.zengde.ai_sandbox.manager;

import com.zengde.ai_sandbox.service.SuperBigModelInterface;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 大模型管理器
 * 负责管理所有会话和历史消息
 */
@Component
@RequiredArgsConstructor
public class BigModelManager {
    
    private final Map<String, SuperBigModelInterface> modelServices;
    
    // 会话缓存: key=sessionId (robotId+customerId), value=modelServiceName
    private final Map<String, String> sessionModelMap = new ConcurrentHashMap<>();
    
    // 机器人角色缓存: key=robotId, value=rolePrompt
    private final Map<String, String> robotRoleMap = new ConcurrentHashMap<>();
    
    /**
     * 创建新的会话
     * @param customerId 客服ID
     * @param modelServiceName 模型服务名称
     * @return 会话ID
     */
    public String createSession(String customerId, String modelServiceName) {
        // 生成机器人ID
        String robotId = generateRobotId();
        
        // 创建会话ID: robotId+customerId
        String sessionId = robotId + "_" + customerId;
        
        // 记录会话使用的模型服务
        sessionModelMap.put(sessionId, modelServiceName);
        
        // 获取角色设定
        // TODO: 从数据库加载角色设定
        String rolePrompt = getRolePromptForRobot(robotId);
        
        // 初始化会话
        SuperBigModelInterface modelService = modelServices.get(modelServiceName);
        modelService.initializeSession(sessionId, rolePrompt);
        
        return sessionId;
    }
    
    /**
     * 非深度思考、非流式调用
     */
    public String ntns(String sessionId, String content) {
        SuperBigModelInterface modelService = getModelServiceForSession(sessionId);
        return modelService.ntns(sessionId, content);
    }
    
    /**
     * 非深度思考、流式调用
     */
    public Flux<String> nts(String sessionId, String content) {
        SuperBigModelInterface modelService = getModelServiceForSession(sessionId);
        return modelService.nts(sessionId, content);
    }
    
    /**
     * 深度思考、流式调用
     */
    public Flux<String> ts(String sessionId, String content) {
        SuperBigModelInterface modelService = getModelServiceForSession(sessionId);
        return modelService.ts(sessionId, content);
    }
    
    /**
     * 深度思考、非流式调用
     */
    public String tns(String sessionId, String content) {
        SuperBigModelInterface modelService = getModelServiceForSession(sessionId);
        return modelService.tns(sessionId, content);
    }
    
    /**
     * 图像聊天
     */
    public String imgChat(String sessionId, String content, String imageUrl) {
        SuperBigModelInterface modelService = getModelServiceForSession(sessionId);
        return modelService.imgChat(sessionId, content, imageUrl);
    }
    
    /**
     * 机器人首次发言
     */
    public String robotFirstMessage(String sessionId) {
        // 获取机器人ID
        String robotId = sessionId.split("_")[0];
        
        // 构建初始消息
        // TODO: 从数据库加载初始消息
        String initialMessage = "你好，我是客户" + robotId + "，请问有什么可以帮助我的吗？";
        
        // 调用模型生成回复
        return ntns(sessionId, initialMessage);
    }
    
    /**
     * 根据会话ID获取对应的模型服务
     */
    private SuperBigModelInterface getModelServiceForSession(String sessionId) {
        String modelServiceName = sessionModelMap.get(sessionId);
        if (modelServiceName == null) {
            throw new IllegalArgumentException("会话不存在: " + sessionId);
        }
        
        SuperBigModelInterface modelService = modelServices.get(modelServiceName);
        if (modelService == null) {
            throw new IllegalArgumentException("模型服务不存在: " + modelServiceName);
        }
        
        return modelService;
    }
    
    /**
     * 获取机器人的角色设定
     */
    private String getRolePromptForRobot(String robotId) {
        // TODO: 从数据库加载角色设定
        // 这里先使用固定的角色设定
        String rolePrompt = "你模拟一个客户，扮演成以下角色，和客服聊天\n" +
                "角色：一个想购买手机的客户，预算在5000元左右，对拍照功能要求较高。";
        
        // 缓存角色设定
        robotRoleMap.put(robotId, rolePrompt);
        
        return rolePrompt;
    }
    
    /**
     * 生成机器人ID
     */
    private String generateRobotId() {
        return "robot_" + UUID.randomUUID().toString().substring(0, 8);
    }
    
    /**
     * 删除会话
     */
    public void removeSession(String sessionId) {
        sessionModelMap.remove(sessionId);
        // TODO: 如果使用Redis，还需要清理Redis中的会话数据
    }
} 